package com.caijg.utils;

import com.alibaba.cloud.ai.advisor.DocumentRetrievalAdvisor;
import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.alibaba.cloud.ai.dashscope.rag.DashScopeDocumentRetriever;
import com.alibaba.cloud.ai.dashscope.rag.DashScopeDocumentRetrieverOptions;
import com.caijg.model.entity.Model;
import com.caijg.service.IModelService;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.stereotype.Component;

import static com.caijg.constant.ModelType.TONGUE_MODEL_TYPE;

/**
 * @author 蔡建国
 * Date  2025/5/10 21:48
 */
@Component
@RequiredArgsConstructor
public class ChatClientHolder {

    private final IModelService modelService;

    // 普通对话模型客户端
    @Getter
    private ChatClient simpleChatClient;
    // 健康评测模型客户端
    @Getter
    private ChatClient evaluateChatClient;
    // 舌苔诊断模型客户端
    @Getter
    private ChatClient tongueChatClient;
    // 测试客户端
    @Getter
    private ChatClient testChatClient;

    // 刷新所有的模型Client
    public void refreshAllChatClient() {
      this.refreshSimpleChatClient();
      this.refreshEvaluateChatClient();
      this.refreshTongueChatClient();
    }

    // 刷新健康评测模型客户端
    private void refreshEvaluateChatClient() {
        // 获取当前选择的对话模型的信息
        Model model = modelService.getCurrentEvaluateModel();
        this.evaluateChatClient = getChatClient(model);
    }

    // 刷新普通对话模型客户端
    private void refreshSimpleChatClient() {
        // 获取当前选择的对话模型的信息
        Model model = modelService.getCurrentSimpleChatModel();
        this.simpleChatClient = getChatClient(model);
    }

    // 刷新普通对话模型客户端
    private void refreshTongueChatClient() {
        // 获取当前选择的对话模型的信息
        Model model = modelService.getcurrentTongueModel();
        this.tongueChatClient = getChatClient(model);
    }

    private ChatClient getChatClient(Model model) {
        DashScopeApi dashscopeApi = new DashScopeApi(model.getApiKey());
        // 知识库名称
        String indexName = model.getKnowledgeBase();

        // 阿里云百炼知识库检索器
        DocumentRetriever retriever = new DashScopeDocumentRetriever(dashscopeApi,
                DashScopeDocumentRetrieverOptions
                        .builder()
                        .withIndexName(indexName)
                        .build());

        return ChatClient.builder(new DashScopeChatModel(dashscopeApi))
                // 添加知识库检索器
                .defaultAdvisors(new DocumentRetrievalAdvisor(retriever))
                .defaultSystem(model.getDefaultSystem())
                .defaultOptions(DashScopeChatOptions
                        .builder()
                        .withModel(model.getModel())
                        .withMultiModel(TONGUE_MODEL_TYPE.equals(model.getType()))
                        .withTemperature(model.getTemperature())
                        .build()
                )
                .build();
    }
}